# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

# -*- coding: utf-8 -*-
"""
测试集合方法的代码生成

验证 set 类型的方法能够正确生成 C++ 代码。
"""

from tests.pytests.helpers.codegen_helpers import compile_chunk, assert_code_contains
from typing import Any


def test_set_add_method():
    """测试 set.add() 方法的代码生成"""

    def test_add(s: set[int]) -> None:
        """测试 add 方法"""
        s.add(42)

    cpp = compile_chunk(test_add)
    expected = """
    void test_add(mcpy::variant_set& s) {
        s.insert(42);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_remove_method():
    """测试 set.remove() 方法的代码生成"""

    def test_remove(s: set[int]) -> None:
        """测试 remove 方法"""
        s.remove(42)

    cpp = compile_chunk(test_remove)
    expected = """
    void test_remove(mcpy::variant_set& s) {
        mcpy::set_remove(s, 42);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_discard_method():
    """测试 set.discard() 方法的代码生成"""

    def test_discard(s: set[int]) -> None:
        """测试 discard 方法"""
        s.discard(42)

    cpp = compile_chunk(test_discard)
    expected = """
    void test_discard(mcpy::variant_set& s) {
        s.erase(42);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_clear_method():
    """测试 set.clear() 方法的代码生成"""

    def test_clear(s: set[int]) -> None:
        """测试 clear 方法"""
        s.clear()

    cpp = compile_chunk(test_clear)
    expected = """
    void test_clear(mcpy::variant_set& s) {
        s.clear();
    }
    """
    assert_code_contains(cpp, expected)


def test_set_pop_method():
    """测试 set.pop() 方法的代码生成"""

    def test_pop(s: set[int]) -> int:
        """测试 pop 方法"""
        return s.pop()

    cpp = compile_chunk(test_pop)
    expected = """
    int test_pop(mcpy::variant_set& s) {
        return mcpy::set_pop<int>(s);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_copy_method():
    """测试 set.copy() 方法的代码生成"""

    def test_copy(s: set[int]) -> set[int]:
        """测试 copy 方法"""
        return s.copy()

    cpp = compile_chunk(test_copy)
    # 引用语义的 set 需要调用 copy() 成员方法创建独立副本
    expected = """
    mcpy::variant_set test_copy(const mcpy::variant_set& s) {
        return s.copy();
    }
    """
    assert_code_contains(cpp, expected)


def test_set_union_method():
    """测试 set.union() 方法的代码生成"""

    def test_union(s1: set[int], s2: set[int]) -> set[int]:
        """测试 union 方法"""
        return s1.union(s2)

    cpp = compile_chunk(test_union)
    expected = """
    mcpy::variant_set test_union(const mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        return mcpy::set_union(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_intersection_method():
    """测试 set.intersection() 方法的代码生成"""

    def test_intersection(s1: set[int], s2: set[int]) -> set[int]:
        """测试 intersection 方法"""
        return s1.intersection(s2)

    cpp = compile_chunk(test_intersection)
    expected = """
    mcpy::variant_set test_intersection(const mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        return mcpy::set_intersection(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_difference_method():
    """测试 set.difference() 方法的代码生成"""

    def test_difference(s1: set[int], s2: set[int]) -> set[int]:
        """测试 difference 方法"""
        return s1.difference(s2)

    cpp = compile_chunk(test_difference)
    expected = """
    mcpy::variant_set test_difference(const mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        return mcpy::set_difference(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_symmetric_difference_method():
    """测试 set.symmetric_difference() 方法的代码生成"""

    def test_symmetric_difference(s1: set[int], s2: set[int]) -> set[int]:
        """测试 symmetric_difference 方法"""
        return s1.symmetric_difference(s2)

    cpp = compile_chunk(test_symmetric_difference)
    expected = """
    mcpy::variant_set test_symmetric_difference(const mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        return mcpy::set_symmetric_difference(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_update_method():
    """测试 set.update() 方法的代码生成"""

    def test_update(s1: set[int], s2: set[int]) -> None:
        """测试 update 方法"""
        s1.update(s2)

    cpp = compile_chunk(test_update)
    expected = """
    void test_update(mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        mcpy::set_update(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_intersection_update_method():
    """测试 set.intersection_update() 方法的代码生成"""

    def test_intersection_update(s1: set[int], s2: set[int]) -> None:
        """测试 intersection_update 方法"""
        s1.intersection_update(s2)

    cpp = compile_chunk(test_intersection_update)
    expected = """
    void test_intersection_update(mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        mcpy::set_intersection_update(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_issubset_method():
    """测试 set.issubset() 方法的代码生成"""

    def test_issubset(s1: set[int], s2: set[int]) -> bool:
        """测试 issubset 方法"""
        return s1.issubset(s2)

    cpp = compile_chunk(test_issubset)
    expected = """
    bool test_issubset(const mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        return mcpy::set_issubset(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_issuperset_method():
    """测试 set.issuperset() 方法的代码生成"""

    def test_issuperset(s1: set[int], s2: set[int]) -> bool:
        """测试 issuperset 方法"""
        return s1.issuperset(s2)

    cpp = compile_chunk(test_issuperset)
    expected = """
    bool test_issuperset(const mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        return mcpy::set_issuperset(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_isdisjoint_method():
    """测试 set.isdisjoint() 方法的代码生成"""

    def test_isdisjoint(s1: set[int], s2: set[int]) -> bool:
        """测试 isdisjoint 方法"""
        return s1.isdisjoint(s2)

    cpp = compile_chunk(test_isdisjoint)
    expected = """
    bool test_isdisjoint(const mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        return mcpy::set_isdisjoint(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_union_operator():
    """测试 set | set 并集运算符的代码生成"""

    def test_union_op(s1: set[int], s2: set[int]) -> set[int]:
        """测试 | 并集运算符"""
        return s1 | s2

    cpp = compile_chunk(test_union_op)
    expected = """
    mcpy::variant_set test_union_op(const mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        return mcpy::set_union(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_intersection_operator():
    """测试 set & set 交集运算符的代码生成"""

    def test_intersection_op(s1: set[int], s2: set[int]) -> set[int]:
        """测试 & 交集运算符"""
        return s1 & s2

    cpp = compile_chunk(test_intersection_op)
    expected = """
    mcpy::variant_set test_intersection_op(const mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        return mcpy::set_intersection(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_difference_operator():
    """测试 set - set 差集运算符的代码生成"""

    def test_difference_op(s1: set[int], s2: set[int]) -> set[int]:
        """测试 - 差集运算符"""
        return s1 - s2

    cpp = compile_chunk(test_difference_op)
    expected = """
    mcpy::variant_set test_difference_op(const mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        return mcpy::set_difference(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_symmetric_difference_operator():
    """测试 set ^ set 对称差集运算符的代码生成"""

    def test_xor_op(s1: set[int], s2: set[int]) -> set[int]:
        """测试 ^ 对称差集运算符"""
        return s1 ^ s2

    cpp = compile_chunk(test_xor_op)
    expected = """
    mcpy::variant_set test_xor_op(const mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        return mcpy::set_symmetric_difference(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_inplace_union_operator():
    """测试 set |= set 原地并集运算符的代码生成"""

    def test_ior_op(s1: set[int], s2: set[int]) -> None:
        """测试 |= 原地并集运算符"""
        s1 |= s2

    cpp = compile_chunk(test_ior_op)
    expected = """
    void test_ior_op(mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        mcpy::set_ior(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_inplace_intersection_operator():
    """测试 set &= set 原地交集运算符的代码生成"""

    def test_iand_op(s1: set[int], s2: set[int]) -> None:
        """测试 &= 原地交集运算符"""
        s1 &= s2

    cpp = compile_chunk(test_iand_op)
    expected = """
    void test_iand_op(mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        mcpy::set_iand(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_inplace_difference_operator():
    """测试 set -= set 原地差集运算符的代码生成"""

    def test_isub_op(s1: set[int], s2: set[int]) -> None:
        """测试 -= 原地差集运算符"""
        s1 -= s2

    cpp = compile_chunk(test_isub_op)
    expected = """
    void test_isub_op(mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        mcpy::set_isub(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_inplace_xor_operator():
    """测试 set ^= set 原地对称差集运算符的代码生成"""

    def test_ixor_op(s1: set[int], s2: set[int]) -> None:
        """测试 ^= 原地对称差集运算符"""
        s1 ^= s2

    cpp = compile_chunk(test_ixor_op)
    expected = """
    void test_ixor_op(mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        mcpy::set_ixor(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_subset_operator():
    """测试 set <= set 子集运算符的代码生成"""

    def test_le_op(s1: set[int], s2: set[int]) -> bool:
        """测试 <= 子集运算符"""
        return s1 <= s2

    cpp = compile_chunk(test_le_op)
    expected = """
    bool test_le_op(const mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        return mcpy::set_issubset(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_superset_operator():
    """测试 set >= set 超集运算符的代码生成"""

    def test_ge_op(s1: set[int], s2: set[int]) -> bool:
        """测试 >= 超集运算符"""
        return s1 >= s2

    cpp = compile_chunk(test_ge_op)
    expected = """
    bool test_ge_op(const mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        return mcpy::set_issuperset(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_proper_subset_operator():
    """测试 set < set 真子集运算符的代码生成"""

    def test_lt_op(s1: set[int], s2: set[int]) -> bool:
        """测试 < 真子集运算符"""
        return s1 < s2

    cpp = compile_chunk(test_lt_op)
    expected = """
    bool test_lt_op(const mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        return mcpy::set_is_proper_subset(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_proper_superset_operator():
    """测试 set > set 真超集运算符的代码生成"""

    def test_gt_op(s1: set[int], s2: set[int]) -> bool:
        """测试 > 真超集运算符"""
        return s1 > s2

    cpp = compile_chunk(test_gt_op)
    expected = """
    bool test_gt_op(const mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        return mcpy::set_is_proper_superset(s1, s2);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_equality_operators():
    """测试 set 的相等和不等运算符"""

    def test_eq_ne(s1: set[int], s2: set[int]) -> bool:
        """测试 == 和 != 运算符"""
        eq: bool = s1 == s2
        ne: bool = s1 != s2
        return eq and ne

    cpp = compile_chunk(test_eq_ne)
    expected = """
    bool test_eq_ne(const mcpy::variant_set& s1, const mcpy::variant_set& s2) {
        bool eq = s1 == s2;
        bool ne = s1 != s2;
        return eq && ne;
    }
    """
    assert_code_contains(cpp, expected)


def test_set_len_function():
    """测试 len(set) 函数的代码生成"""

    def test_len(s: set[int]) -> int:
        """测试 len 函数"""
        return len(s)

    cpp = compile_chunk(test_len)
    expected = """
    int test_len(const mcpy::variant_set& s) {
        return s.size();
    }
    """
    assert_code_contains(cpp, expected)


def test_set_contains_operator():
    """测试 in 运算符的代码生成"""

    def test_contains(s: set[int], item: int) -> bool:
        """测试 in 运算符"""
        return item in s

    cpp = compile_chunk(test_contains)
    expected = """
    bool test_contains(const mcpy::variant_set& s, int item) {
        return mcpy::set_contains(s, item);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_bool_operator():
    """测试 bool(set) 的代码生成"""

    def test_bool(s: set[int]) -> bool:
        """测试 bool 转换"""
        if s:
            return True
        return False

    cpp = compile_chunk(test_bool)
    expected = """
    bool test_bool(const mcpy::variant_set& s) {
        if (!s.empty()) {
            return true;
        }
        return false;
    }
    """
    assert_code_contains(cpp, expected)


def test_set_with_strings():
    """测试字符串类型的 set"""

    def test_string_set(s: set[str]) -> None:
        """测试字符串 set"""
        s.add("hello")
        s.remove("world")

    cpp = compile_chunk(test_string_set)
    expected = """
    void test_string_set(mcpy::variant_set& s) {
        s.insert("hello");
        mcpy::set_remove(s, "world");
    }
    """
    assert_code_contains(cpp, expected)


def test_set_combined_operations():
    """测试组合使用多个 set 运算符"""

    def test_combined(s1: set[int], s2: set[int], s3: set[int]) -> set[int]:
        """组合运算"""
        # (s1 | s2) & s3
        temp: set[int] = s1 | s2
        result: set[int] = temp & s3
        return result

    cpp = compile_chunk(test_combined)
    expected = """
        mcpy::variant_set test_combined(const mcpy::variant_set& s1, const mcpy::variant_set& s2, const mcpy::variant_set& s3) {
            mcpy::variant_set temp   = mcpy::set_union(s1, s2);
            mcpy::variant_set result = mcpy::set_intersection(temp, s3);
            return result;
        }
    """
    assert_code_contains(cpp, expected)


# ============================================================================
# set 构造函数测试
# ============================================================================


def test_set_empty_constructor():
    """测试空 set 构造函数"""

    def test_empty() -> None:
        """测试空构造"""
        s1: set[int] = set()
        s2: set[str] = set()

    cpp = compile_chunk(test_empty)
    expected = """
    void test_empty() {
        mcpy::variant_set         s1 = mcpy::variant_set{};
        mcpy::variant_set s2 = mcpy::variant_set{};
    }
    """
    assert_code_contains(cpp, expected)


def test_set_from_list_constructor():
    """测试从 list 构造 set"""

    def test_from_list() -> None:
        """从列表构造 set"""
        numbers: list[int] = [1, 2, 2, 3, 3, 3]
        s: set[int] = set(numbers)

    cpp = compile_chunk(test_from_list)
    expected = """
    void test_from_list() {
        mc::array<int> numbers = {1, 2, 2, 3, 3, 3};
        mcpy::variant_set   s       = mcpy::set_construct<int>(numbers);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_from_list_literal_constructor():
    """测试从列表字面量构造 set"""

    def test_from_literal() -> None:
        """从列表字面量构造 set"""
        s: set[int] = set([1, 2, 3, 4, 5])

    cpp = compile_chunk(test_from_literal)
    expected = """
    void test_from_literal() {
        mcpy::variant_set s = mcpy::set_construct<int>(mc::array<int>{1, 2, 3, 4, 5});
    }
    """
    assert_code_contains(cpp, expected)


def test_set_from_set_constructor():
    """测试从另一个 set 拷贝构造"""

    def test_copy_constructor(s1: set[int]) -> set[int]:
        """拷贝构造 set"""
        s2: set[int] = set(s1)
        return s2

    cpp = compile_chunk(test_copy_constructor)
    expected = """
    mcpy::variant_set test_copy_constructor(const mcpy::variant_set& s1) {
        mcpy::variant_set s2 = mcpy::set_construct<int>(s1);
        return s2;
    }
    """
    assert_code_contains(cpp, expected)


def test_set_constructor_type_inference():
    """测试 set 构造函数的类型推导"""

    def test_inference() -> None:
        """测试类型推导"""
        # 带类型注解
        s1: set[str] = set(["a", "b", "c"])
        # 从已知类型列表推导
        numbers: list[int] = [1, 2, 3]
        s2: set[int] = set(numbers)

    cpp = compile_chunk(test_inference)
    expected = """
    void test_inference() {
        mcpy::variant_set s1 = mcpy::set_construct<std::string>(mc::array<std::string>{"a", "b", "c"});
        mc::array<int>       numbers = {1, 2, 3};
        mcpy::variant_set         s2      = mcpy::set_construct<int>(numbers);
    }
    """
    assert_code_contains(cpp, expected)


def test_set_constructor_in_expression():
    """测试 set 构造函数在表达式中使用"""

    def test_in_expr() -> int:
        """在表达式中使用 set 构造"""
        return len(set([1, 2, 2, 3, 3, 3]))

    cpp = compile_chunk(test_in_expr)
    expected = """
    int test_in_expr() {
        return mcpy::set_construct<int>(mc::array<int>{1, 2, 2, 3, 3, 3}).size();
    }
    """
    assert_code_contains(cpp, expected)


# ============================================================================
# 弱类型 set 测试（无类型参数）
# ============================================================================


def test_weak_type_set_creation():
    """测试弱类型 set 的创建"""

    def test_weak_set() -> None:
        """测试弱类型 set"""
        # 无类型注解的 set
        s = {1, "hello", 3.14}

    cpp = compile_chunk(test_weak_set)
    expected = """
    void test_weak_set() {
        mcpy::variant_set s = {1, "hello", 3.14};
    }
    """
    assert_code_contains(cpp, expected)


def test_weak_type_set_operations():
    """测试弱类型 set 的操作"""

    def test_operations() -> int:
        """测试弱类型 set 的操作"""
        s: set = {1, 2, 3}
        s.add(4)
        s.remove(2)
        return len(s)

    cpp = compile_chunk(test_operations)
    expected = """
    int test_operations() {
        mcpy::variant_set s = {1, 2, 3};
        s.insert(4);
        mcpy::set_remove(s, 2);
        return s.size();
    }
    """
    assert_code_contains(cpp, expected)


def test_weak_type_set_mixed_types():
    """测试弱类型 set 混合类型"""

    def test_mixed() -> bool:
        """测试混合类型 set"""
        s: set = {1, "hello", True, 3.14}
        return 1 in s and "hello" in s

    cpp = compile_chunk(test_mixed)
    expected = """
    bool test_mixed() {
        mcpy::variant_set s = {1, "hello", true, 3.14};
        return mcpy::set_contains(s, 1) && mcpy::set_contains(s, "hello");
    }
    """
    assert_code_contains(cpp, expected)


def test_weak_type_set_union():
    """测试弱类型 set 的并集运算"""

    def test_union() -> int:
        """测试弱类型并集"""
        s1: set = {1, 2}
        s2: set = {2, 3}
        s3: set = s1 | s2  # 需要类型注解
        return len(s3)

    cpp = compile_chunk(test_union)
    expected = """
    int test_union() {
        mcpy::variant_set s1 = {1, 2};
        mcpy::variant_set s2 = {2, 3};
        mcpy::variant_set s3 = mcpy::set_union(s1, s2);
        return s3.size();
    }
    """
    assert_code_contains(cpp, expected)
